// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <random>

#include <pollux/exec/exchange_client.h>
#include <pollux/exec/operator.h>
#include <pollux/exec/output_buffer_manager.h>
#include <pollux/serializers/presto_serializer.h>

namespace kumo::pollux::exec {
    struct RemoteConnectorSplit : public connector::ConnectorSplit {
        const std::string taskId;

        explicit RemoteConnectorSplit(const std::string &remoteTaskId)
            : ConnectorSplit(""), taskId(remoteTaskId) {
        }

        std::string toString() const override {
            return fmt::format("Remote: {}", taskId);
        }
    };

    class Exchange : public SourceOperator {
    public:
        Exchange(
            int32_t operatorId,
            DriverCtx *driverCtx,
            const std::shared_ptr<const core::ExchangeNode> &exchangeNode,
            std::shared_ptr<ExchangeClient> exchangeClient,
            const std::string &operatorType = "Exchange");

        ~Exchange() override {
            close();
        }

        RowVectorPtr getOutput() override;

        void close() override;

        BlockingReason isBlocked(ContinueFuture *future) override;

        bool isFinished() override;

    protected:
        virtual VectorSerde *getSerde();

    private:
        // Invoked to create exchange client for remote tasks. The function shuffles
        // the source task ids first to randomize the source tasks we fetch data from.
        // This helps to avoid different tasks fetching from the same source task in a
        // distributed system.
        void addRemoteTaskIds(std::vector<std::string> &remoteTaskIds);

        // Fetches splits from the task until there are no more splits or task returns
        // a future that will be complete when more splits arrive. Adds splits to
        // exchangeClient_. Returns true if received a future from the task and sets
        // the 'future' parameter. Returns false if fetched all splits or if this
        // operator is not the first operator in the pipeline and therefore is not
        // responsible for fetching splits and adding them to the exchangeClient_.
        bool getSplits(ContinueFuture *future);

        // Fetches runtime stats from ExchangeClient and replaces these in this
        // operator's stats.
        void recordExchangeClientStats();

        const uint64_t preferredOutputBatchBytes_;

        const VectorSerde::Kind serdeKind_;

        const std::unique_ptr<VectorSerde::Options> serdeOptions_;

        /// True if this operator is responsible for fetching splits from the Task
        /// and passing these to ExchangeClient.
        const bool processSplits_;

        const int driverId_;

        bool noMoreSplits_ = false;

        std::shared_ptr<ExchangeClient> exchangeClient_;

        // A future received from Task::getSplitOrFuture(). It will be complete when
        // there are more splits available or no-more-splits signal has arrived.
        ContinueFuture splitFuture_{ContinueFuture::makeEmpty()};

        // Reusable result vector.
        RowVectorPtr result_;

        std::vector<std::unique_ptr<SerializedPage> > currentPages_;
        bool atEnd_{false};
        std::default_random_engine rng_{std::random_device{}()};
    };
} // namespace kumo::pollux::exec
